"""
输入内容
时间：2024/8/22 下午2:50
"""
import warnings
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from flask import Flask, request, jsonify


# 创建一个Flask应用实例
app = Flask(__name__)

# 忽略所有警告
warnings.filterwarnings("ignore")


# 计算相似度
def calculate_similarity(text1, text2):
    # pipeline加载模型
    similarity_pipeline = pipeline(
        task=Tasks.sentence_similarity,
        model='damo/nlp_structbert_sentence-similarity_chinese-retail-base',
        model_revision='v1.0.0'
    )
    # 得到模型输出
    pipeline_result = similarity_pipeline(input=(text1, text2))
    # print(pipeline_result)
    # 找到标签为‘相似’的标签值
    similarity_index = pipeline_result['labels'].index('相似')
    # 计算相似度百分比
    similarity_percentage = pipeline_result['scores'][similarity_index] * 100
    # print(f'similarity:{similarity_percentage:.2f}%')
    return similarity_percentage


# 定义了一个路由，指向应用程序的根 URL
@app.route('/calculate_similarity', methods=['POST'])
def calculate_similarity_endpoint():
    # 拿到请求数据
    data = request.get_json()
    if not data or 'text' not in data:
        return jsonify({"error": "Missing required field 'text'"}), 400
    # 加载文本
    text1 = data.get('text_1')          # 要保证数据的正确性
    text2 = data.get('text_2')
    # 计算相似度
    result = calculate_similarity(text1, text2)

    response = {
        'input': {
            'input1': text1,
            'input2': text2
        },
        'output':{
            'similarity': f'{result:.2f}%'
        }
    }

    return jsonify(response), 200


if __name__ == '__main__':
    app.run(debug=True, port=40006)
